function [x_opt,a_opt] = optimization_NAO
global ns_time s_time ns_slope s_knee ns_knee hip_pos pos_com

%%
%%%%%%%%%%%%%Add paths to mathematica compiled files:
close all
addpath('./build_torso')
addpath('./wrappers_torso')
addpath('./buildopt_torso')
addpath('./data')

%%
%%%%%%%%%%%%% Loading the data:
%%%NOTE:  Make sure to pick the data to match the outputs chosen

%%%%%%%%%%%%%Loading specific data:
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%% Hip position
hipdatatype = 2;
if hipdatatype == 1
    load('mean_data_hippos.mat')
    %     hip_pos = hip_pos;
elseif hipdatatype == 2
    load('mean_data_LinearHippos.mat')
    hip_pos = Linear_hippos;
elseif hipdatatype == 3
    load('mean_data_LinearHipposD.mat')
    hip_pos = Linear_hippos_D;
elseif hipdatatype == 4
    load('mean_data_LinearHippos_shift.mat')
    hip_pos = Linear_nsslope_shift;
end
% ns_time = ns_time;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%% Non-stance slope or hip angle
nsslopedatatype = 2;
if nsslopedatatype == 1
    load('mean_data_nsslope.mat')
    % ns_slope = ns_slope
elseif nsslopedatatype == 2
    load('mean_data_LinearNSslope.mat')
    ns_slope = Linear_nsslope;
elseif nsslopedatatype == 3
    load('mean_data_hipAngle.mat')
    ns_slope = hip_angle;
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%% Stance knee
load('mean_data_sknee.mat')
% s_knee = s_knee
s_time = ns_time ;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%% Non-stance knee
load('mean_data_nsknee.mat')
% ns_knee = ns_knee;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%% Torso
torsodatatype = 2;
if torsodatatype == 1
    load('mean_data_TorsoHipAngle.mat')
    pos_com = torso_angle;
elseif torsodatatype == 2
    load('mean_data_theta4.mat')
    pos_com = theta4;
elseif torsodatatype == 3
    load('mean_data_nstorso.mat')
    pos_com = ns_torso;
elseif torsodatatype == 4
    load('mean_data_LinearNStorso.mat')
    pos_com = Linear_ns_torso;
elseif torsodatatype == 5
    load('mean_data_storso.mat')
    pos_com = s_torso;
end

%%%%%%%%%%%%%Initial fit parameters (as determined my Mathematica, in the
%%%%%%%%%%%%%Data folder.  Note, make sure that the fit parameters match
%%%%%%%%%%%%%the outputs chose.

load('Afit.mat', 'a0');

%%%%%%%%%%%%%% Minimization to verify optimality of initial parameters w.r.t the HBC
% options = optimset ('Display','iter','MaxFunEvals',6000);
% [a_fit,cost_fit] = fminunc(@cost,flatten(a0),options)
%%%%%%%%%%%%%% In the case that you don't want to run the optimizatino to
%%%%%%%%%%%%%% verify the parameters, uncomment this:
a_fit = flatten(a0);

%%%%%%%%%%%%%% The matrix constrains the values of the parameters
%%%don't want the paramters to get too large, casues problems with
%%%integration
ubmat = 50*ones(5,5);
lbmat = -50*ones(5,5);
%%upper and lower bounds on the hip velocity of the robot
ubmat(1,1) = 0.8; %%% SHU
lbmat(1,1) = 0.1;

%%Add any constraints to the torso fit function here
% ubmat(5,:) = 0.1*ones(1,5);
% lbmat(5,:) = -0.0*ones(1,5);
%%%
ubmat(1,3) = 0; ubmat(1,4) = 0; ubmat(1,5) = 0;
lbmat(1,3) = 0; lbmat(1,4) = 0; lbmat(1,5) = 0;
if hipdatatype ~= 4
    lbmat(1,2) = 0
    ubmat(1,2) = 0
end
%%%
ub = flatten(ubmat);
lb = flatten(lbmat);

%%%%%%%%%%%%%% Other conditions for the optimization:
Aeq = [];
beq = [];
A = [];
b = [];

%%%%%%%%%%%%%% Optimizing to achieve HZD with the HBC:

%%%%% Conditions for optimization: tight error
errorbound = 1e-8;
options = optimset('Algorithm','sqp',...
    'Display','iter','TolX',errorbound,...
    'MaxFunEvals',20000,'TolCon',errorbound);

%%%%% Optimization
tic
[a_out,cost_out] = fmincon(@cost,a_fit,A,b,Aeq,beq,lb,ub,@mycon,options);
toc

%%%%%%%%%%%%%% Output the parameter matrix and the fit matrix:
a_opt = unflatten(a_out);
a_fit = unflatten(a_fit);

%%%%%%%%%%%%%% Output the initial condition:
[x_opt,tau_opt] = hzd_con(a_opt);

%%%%%%%%%%%%%% Output the final time:
disp(['Step time: ' num2str(tau_opt)]);

%%%%%%%%%%%%%% Save initial and final condition
save('x_opt.mat','x_opt')
save('a_opt.mat','a_opt')

%%%%%%%%%%%%%% Plot the output vs. fit vs. the human data
fig12=figure(12); clf;
subplot(2,3,1); plot(ns_time,yd_hip_pos(ns_time,a_opt),'b',ns_time,yd_hip_pos(ns_time,a_fit),'r',ns_time,hip_pos,'k.')
title('Hip Position')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,2); plot(ns_time,yd_ns_slope(ns_time,a_opt),'b',ns_time,yd_ns_slope(ns_time,a_fit),'r',ns_time,ns_slope,'k.')
title('Non-Stance Slope')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,4); plot(s_time,yd_s_knee(s_time,a_opt),'b',s_time,yd_s_knee(s_time,a_fit),'r',s_time,s_knee,'k.')
title('Stance Knee')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,5); plot(ns_time,yd_ns_knee(ns_time,a_opt),'b',ns_time,yd_ns_knee(ns_time,a_fit),'r',ns_time,ns_knee,'k.')
title('Non-Stance Knee')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,3); plot(ns_time,yd_pos_com(ns_time,a_opt),'b',ns_time,yd_pos_com(ns_time,a_fit),'r',ns_time,pos_com,'k.')
title('Torso')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
figname = strcat('result/test.pdf');
% saveas(fig12,'LNST4.pdf')
saveas(fig12,figname)
%%%%%%%%%%%%%% Run main
main_NAO('Test')

end

%%
%%%%%%%%%%%%%% Cost function (based upon human data)

function ret = cost(aflat)
global ns_time s_time ns_slope s_knee ns_knee hip_pos pos_com
global typeofcost
global kneemin hmin hmax qdotmax umax deltaP powermax kneemax

%%%%%%%%%%%%%% Unflatten the a matrix
a = unflatten(aflat);

%%%%%Pick the cost function
%%% 0 = torque based cost, requires additional integration
%%% 1 = human data based cost
typeofcost = 1;

%%%%%The torque based cost
if typeofcost == 0
    
    %%% Determine the point on the guard through inverse kinematics
    [x_minus,Ttau,yminus,H,phipcond] = hzd_con(a);
    %%% Integrate the system forward one step
    x_plus = resetFunc(x_minus);
    [Tf,xf] = step_under_reduced(x_plus,a,Ttau);
    
    %%% Determine the step length (if needed
    posf = jpos_mat(xf);
    steplength = posf(1,end);
    if steplength < 0
        steplength = 0.0001;
    end
    
    %%% Compute the cost
    % ret =  norm(xf-x_minus)^2;
    % ret = umax/abs(steplength)+norm(xf-x_minus)^2;
    % ret = umax/abs(steplength);
    %  ret = umax/steplength;
    ret = (umax)/steplength;
    
    %%%% The data based cost
elseif typeofcost == 1
    
    %%%% Decide if you want to scale time based upon the parameterization of
    %%%% time
    scaletime = 0;
    if scaletime == 0
        
        time_scale = ns_time;
        
    elseif scaletime == 1
        
        [x_minus,Ttau,yminus,H,phipcond] = hzd_con(a);
        time_scale = ns_time*(Ttau/ns_time(end));
        
    end
    
    %%%%%%%%%%%%%% Cost for each canonical function
    
    cost_ns_slope = sum(abs(ns_slope-yd_ns_slope(time_scale,a)));
    cost_hip_pos = sum(abs(hip_pos-yd_hip_pos(time_scale,a)));
    cost_s_knee = sum(abs(s_knee-yd_s_knee(time_scale,a)));
    cost_ns_knee = sum(abs(ns_knee-yd_ns_knee(time_scale,a)));
    cost_pos_com = sum(abs(pos_com-yd_pos_com(time_scale,a)));
    
    %%%%%%%%%%%%%% Wieghting for the fits (in this case the torso cost is
    %%%%%%%%%%%%%% not used.
   
    beta_ns_slope = 1/(max(ns_slope)-min(ns_slope));
    beta_hip_pos = 1/(max(hip_pos)-min(hip_pos));
    beta_s_knee = 1/(max(s_knee)-min(s_knee));
    beta_ns_knee = 1/(max(ns_knee)-min(ns_knee));
    beta = beta_ns_slope+beta_hip_pos+beta_s_knee+beta_ns_knee;
% %         beta_pos_com = 1; %/(max(pos_com)-min(pos_com));
    beta_pos_com = 1/(max(pos_com)-min(pos_com));
        beta = beta_ns_slope+beta_hip_pos+beta_s_knee+beta_ns_knee+beta_pos_com;
    
    %%% Determine the cost
% %     cost_all = [
% %         beta_ns_slope*cost_ns_slope,...
% %         beta_hip_pos*cost_hip_pos,...
% %         beta_s_knee*cost_s_knee,...
% %         beta_ns_knee*cost_ns_knee
% %         ];
% %     
        cost_all = [
              beta_ns_slope*cost_ns_slope,...
              beta_hip_pos*cost_hip_pos,...
              beta_s_knee*cost_s_knee,...
              beta_ns_knee*cost_ns_knee,...
              beta_pos_com*cost_pos_com
              ];
    
    %%%%%%%%%%%%%% Return the cost
    ret = sum(cost_all)/beta;
end

end


%%
%%%%%%%%%%%%%% Constraints to enforce HZD (or partial HZD)

function [c,ceq] = mycon(aflat)
global ns_time 

%%%%%%%%%%%%%% Unflatten the a matrix
a = unflatten(aflat);

%%
%%%%Read in the constraints for HZD
[x_minus,Ttau,yminus,H] = hzd_con(a);
x_plus = resetFunc(x_minus);

%%
%%%%%%%%%%%%%% Require the post impact velocities be completlty determined:
%%% partial hybrid zero dynamics
% Hzd = H(2:end,:) - [0; 0; 0; 0];
%%% full hybrid zero dynamics
Hzd = H - [a(1,1); 0; 0; 0; 0];

%%%%%%%%%%%%%% Pre and Post Impact knee angles

%%%%%%%%%%%%%% Nonlinear inequalities
%%% This is where you specify any constraints on the system
dimq = length(x_minus)/2;

    %%%%%%%%%%%%%% Nonlinear inequalities
    %%%%%%%%%%%%%%%%% gaurd must be negative
    %%%%%%%%%%%%%%%%% Time must be less than the human time
    c = [
           h_dot_mat(x_minus(1:dimq))*x_minus(dimq+1:end)+.01;
        ];
    
%%%%%%%%%%%%%% Nonlinear equalities
%%%%%%%%%%%%%%%%% y pre impact must be zero
%%%%%%%%%%%%%%%%% y dot post impact must be zero
%%%%%%%%%%%%%%%%% guard must be satisfied
ceq = [
       yminus;
       Ttau-ns_time(end);
%        a(1,2) - deltaphip_sca(x_plus);
       Hzd];

end


function [x_minus,Ttau,yminus,H,phipcond] = hzd_con(a)

%%%%%%%%%%Determine the initial position on the guard

q_minus = inverse_kin(a);

%%
%%%%%%%%%%Determine the pre and post impact values of the output functions

%%%%Initial position of the hip
dimq = length(q_minus);
q_plus = R(q_minus);
phipcond = phip_sca(q_plus);

%%%%Pre-impact relative degree 2 outputs
yminus = ya2_vec(q_minus, phipcond, a)-yd2_vec(q_minus, phipcond,  a);
% y2plus = ya2_vec(q_plus, phipcond, a)-yd2_vec(q_plus, phipcond,  a)

%%%%Pre-and post impact relative Jacobians of the outputs

Dy2_plus = Dya2_mat(q_plus, phipcond, a)-Dyd2_mat(q_plus, phipcond,  a);
Hplus = [ phip_dot_mat(q_plus);
    Dy2_plus(:,1:dimq)];

Dy2_minus = Dya2_mat(q_minus, phipcond, a)-Dyd2_mat(q_minus, phipcond,  a);
Hminus = [ phip_dot_mat(q_minus);
    Dy2_minus(:,1:dimq)];

%%%%%%%%%%The Time to impact (assuming no scuffing)
Ttau = sigma_sca(q_minus,phipcond,a);

%%%%Pre- and Post-impact velocities

qdot_minus = Hminus\[a(1,1); 0; 0; 0; 0];
x_minus = [q_minus; qdot_minus];

H = Hplus*P(x_minus);

end

%%
%%%%%%%Computing the position at impact through inverse kinematics

function ret = inverse_kin(a)

%%%%Choose the type of inverse kinematics, either numeric or closed form
%%% 1 = numeric solving
%%% 2 = closed form solution (currently only works for non-stance slope and
%%% linearized non-stance slope)
typeofsolve = 1;

if typeofsolve == 1
    
    %%%%% Pick initial guess
    %%%%%q0 as computed from the human parameters
    
    q0 = [ -0.5138
        0.2776
        0.2870
        -0.4201
        0.1746];
    
    %     q0 =  [   -0.4248
    %     0.3006
    %     0.1242
    %    -0.3693
    %     0.1416];
    
    %     load('x_opt')
    %     q0 = x_opt(1:5,:);
    
    %%%Numerically solving for the initial condition
    
    options = optimset('Display','off');
    qzero = fsolve(@(q)inverse_kin_fun(q,a),q0,options);
    
    ret = qzero;
    
elseif typeofsolve == 2
    
    %%%%If using the old version (no torso) use this:
    
    % theta_a(a)
    
    %%%If using the new version with the torso, use this:
    %%%%(the calculation is broken up, saves a lot of computation time)
    
    q = zeros(5,1);
    q(2) = theta_a2(q,a);
    q(5) = theta_a5(q,a);
    q(1) = theta_a1(q,a);
    q(4) = theta_a4(q,a);
    q(3) = theta_a3(q,a);
    
    ret = q;
    
end


end

function F = inverse_kin_fun(q,a)

phipcond = phip_sca(R(q));
y2plus = ya2_vec(R(q), phipcond, a)-yd2_vec(R(q), phipcond,  a);

F = [y2plus;
    h_sca(q)];

end

%%
%%%%%%%%%%%%%% Computing the position and velocity components of the impact
%%%%%%%%%%%%%% equations

function ret = R(q)

dimq = length(q);
x = [q; zeros(dimq,1)];
xplus = resetFunc(x);

ret = xplus(1:dimq);

end

function ret = P(x)

dimq = length(x)/2;
xplus = resetFunc(x);

ret = xplus(dimq+1:end);

end

%%
%%%%%%%%%%%%%% Canonical human functions

function ret = yd_ns_slope(t,a)

ret =exp(-a(2,4)*t).*(a(2,1)*cos(a(2,2)*t)+a(2,3)*sin(a(2,2)*t))+a(2,5);

end

function ret = yd_hip_pos(t,a)

ret = a(1,1)*t+a(1,2);

end

function ret = yd_s_knee(t,a)

ret = exp(-a(3,4)*t).*(a(3,1)*cos(a(3,2)*t)+a(3,3)*sin(a(3,2)*t))+a(3,5);

end

function ret = yd_ns_knee(t,a)

ret = exp(-a(4,4)*t).*(a(4,1)*cos(a(4,2)*t)+a(4,3)*sin(a(4,2)*t))+a(4,5);

end

function ret = yd_pos_com(t,a)

ret = exp(-a(5,4)*t).*(a(5,1)*cos(a(5,2)*t)+a(5,3)*sin(a(5,2)*t))+a(5,5);

end

%%
%%%%%%%%%%%%%% Flatten and unflatten the a matrix

function aflat = flatten(a)
aflat = a(:);

end

function aunflat = unflatten(aflat)

qdim = 5;
aunflat = reshape(aflat,qdim,5);

end


